package asdf.utils;

import org.apache.commons.codec.binary.Base64;

import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import java.io.ByteArrayOutputStream;
import java.security.*;
import java.security.interfaces.RSAPrivateKey;
import java.security.interfaces.RSAPublicKey;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;

/**
 * Created by lichengwen on 2016/10/26.
 */
public class CodecUtils
{
    
    /**
     * decode with base64
     */
    public static byte[] toBytes(String s)
    {
        return Base64.decodeBase64(s);
    }
    
    /**
     * encode with base64
     */
    public static String toString(byte[] b)
    {
        return Base64.encodeBase64String(b);
    }
    
    private static PrivateKey buildPrivateKey(String sPrivateKey, String keyAlgorithm)
    {
        try {
            return KeyFactory.getInstance(keyAlgorithm).generatePrivate(
                    new PKCS8EncodedKeySpec(toBytes(sPrivateKey)));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
    
    private static PublicKey buildPublicKey(String sPublicKey, String keyAlgorithm)
    {
        try {
            return KeyFactory.getInstance(keyAlgorithm).generatePublic(
                    new X509EncodedKeySpec(toBytes(sPublicKey)));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
    
    /**
     * 用私钥对信息生成数字签名
     */
    public static String sign(byte[] data, String privateKey, String keyAlgorithm, String signatureAlgorithm)
    {
        try {
            Signature signature = Signature.getInstance(signatureAlgorithm);
            signature.initSign(buildPrivateKey(privateKey, keyAlgorithm));
            signature.update(data);
            return toString(signature.sign());
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
    
    /**
     * 校验数字签名
     */
    public static boolean checkSign(
            byte[] data, String publicKey, String sign, String keyAlgorithm, String signatureAlgorithm)
    {
        try {
            Signature signature = Signature.getInstance(signatureAlgorithm);
            signature.initVerify(buildPublicKey(publicKey, keyAlgorithm));
            signature.update(data);
            return signature.verify(toBytes(sign));
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
    
    /**
     * 解密
     *
     * @param isPrivateKey true for PrivateKey, false for PublicKey
     */
    public static byte[] decrypt(byte[] data, String key, boolean isPrivateKey, String keyAlgorithm)
    {
        try {
            Cipher cipher = Cipher.getInstance(keyAlgorithm);
            cipher.init(Cipher.DECRYPT_MODE,
                        (isPrivateKey ? buildPrivateKey(key, keyAlgorithm) : buildPublicKey(key, keyAlgorithm)));
            return doCipher(data, 128, cipher);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
    
    /**
     * 加密
     *
     * @param isPublicKey true for PublicKey, false for PrivateKey
     */
    public static byte[] encrypt(byte[] data, String key, boolean isPublicKey, String keyAlgorithm)
    {
        try {
            Cipher cipher = Cipher.getInstance(keyAlgorithm);
            cipher.init(Cipher.ENCRYPT_MODE,
                        (isPublicKey ? buildPublicKey(key, keyAlgorithm) : buildPrivateKey(key, keyAlgorithm)));
            return doCipher(data, 117, cipher);
        } catch (Exception e) {
            throw new RuntimeException(e);
        }
    }
    
    private static byte[] doCipher(byte[] data, int maxLen, Cipher cipher)
            throws IllegalBlockSizeException, BadPaddingException
    {
        if (data.length <= maxLen) {
            return cipher.doFinal(data);
        } else {
            ByteArrayOutputStream byteStream = new ByteArrayOutputStream();
            for (int i = 0; i < data.length; i = i + maxLen) {
                byte[] snippet;
                if (data.length < i + maxLen) {
                    snippet = new byte[data.length - i];
                } else {
                    snippet = new byte[maxLen];
                }
                System.arraycopy(data, i, snippet, 0, snippet.length);
                byte[] encryptedSnippet =  cipher.doFinal(snippet);
                byteStream.write(encryptedSnippet, 0, encryptedSnippet.length);
            }
            return byteStream.toByteArray();
        }
    }
    
    public static void test() throws NoSuchAlgorithmException
    {
        KeyPairGenerator keyPairGen = KeyPairGenerator.getInstance("RSA");
        keyPairGen.initialize(1024);
        KeyPair keyPair = keyPairGen.generateKeyPair();
        RSAPublicKey publicKey = (RSAPublicKey) keyPair.getPublic();
        RSAPrivateKey privateKey = (RSAPrivateKey) keyPair.getPrivate();
        String pubKey = Base64.encodeBase64String(publicKey.getEncoded());
        String priKey = Base64.encodeBase64String(privateKey.getEncoded());
        System.out.println(pubKey);
        System.out.println(priKey);
        
        String s = "asdf阿瓦kj1234567980";
        System.out.println(s);
        byte[] b = s.getBytes();
        
        byte[] eb = encrypt(b, pubKey, true, "RSA");
        System.out.println(new String(eb));
        byte[] db = decrypt(eb, priKey, true, "RSA");
        System.out.println(new String(db));
        
        eb = encrypt(b, priKey, false, "RSA");
        System.out.println(new String(eb));
        db = decrypt(eb, pubKey, false, "RSA");
        System.out.println(new String(db));
    
        String sign = sign(b, priKey, "rsa", "MD5withRSA");
        System.out.println("sign:" + sign);
        System.out.println(checkSign(b, pubKey, sign, "rsa", "MD5withRSA"));
    }
    
    public static void printAlgorithms()
    {
        for (Provider provider : Security.getProviders()) {
            System.out.println("Provider: " + provider.getName());
            for (Provider.Service service : provider.getServices()) {
                System.out.println("  Algorithm: " + service.getAlgorithm());
            }
        }
        
    }
    
    public static void main(String[] args) throws NoSuchAlgorithmException
    {
        test();
    }
}
